# -*- coding:utf-8 -*-
import os  # 导入操作系统模块，用于文件路径操作
from typing import Union, List  # 从 typing 模块导入 Union 和 List，用于类型提示
# Union 是 typing 模块中定义的一个类，用于表示多个类型中的任意一种类型
from pet_config import *  # 导入配置文件，包含项目相关参数
pc = ProjectConfig()  # 实例化项目配置类


class Verbalizer(object):
    """
    Verbalizer类，用于将一个Label对应到其子Label的映射。
    """

    def __init__(self, verbalizer_file: str, tokenizer, max_label_len: int):
        """
        初始化Verbalizer类。

        Args:
            verbalizer_file (str): verbalizer文件存放地址。
            tokenizer: 用于文本和id之间的转换。
            max_label_len (int): 标签长度，若大于则截断，若小于则补齐。
        """
        self.tokenizer = tokenizer  # 存储分词器实例
        self.label_dict = self.load_label_dict(verbalizer_file)  # 加载标签字典
        self.max_label_len = max_label_len  # 设置最大标签长度

    def load_label_dict(self, verbalizer_file: str):
        """
        读取本地文件，构建verbalizer字典。

        Args:
            verbalizer_file (str): verbalizer文件存放地址。

        Returns:
            dict -> {
                '体育': ['篮球', '足球','网球', '排球',  ...],
                '酒店': ['宾馆', '旅馆', '旅店', '酒店', ...],
                ...
            }
        """
        label_dict = {}  # 初始化标签字典
        with open(verbalizer_file, 'r', encoding='utf8') as f:  # 打开文件并读取内容
            for line_num, line in enumerate(f.readlines(), start=1):  # 遍历文件的每一行
                line = line.strip()  # 去除首尾空白字符
                if not line:  # 跳过空行
                    print(f"[Warning] Line {line_num} is empty.")
                    continue
                print(f"Raw line {line_num}: '{line}'")  # 打印原始行内容
                try:
                    parts = line.split('\t')  # 按制表符分割主标签和子标签
                    print(f"Parsed parts: {parts}")  # 打印分割结果
                    if len(parts) != 2:  # 如果分割结果不是两部分，抛出异常
                        raise ValueError(f"Line {line_num} does not have exactly two parts.")
                    label, sub_labels = parts  # 分别获取主标签和子标签
                    label_dict[label] = list(set(sub_labels.split(',')))  # 将子标签去重后存储为列表
                except ValueError as e:  # 捕获异常并打印错误信息
                    print(f"[Error] Line {line_num} is malformed: {e}")
        return label_dict  # 返回构建好的标签字典

    def find_sub_labels(self, label: Union[list, str]):
        """
        通过标签找到所有的子标签。

        Args:
            label (Union[list, str]): 标签, 文本型 或 id_list, e.g. -> '体育' or [860, 5509]

        Returns:
            dict -> {
                'sub_labels': ['足球', '网球'],
                'token_ids': [[6639, 4413], [5381, 4413]]
            }
        """
        if type(label) == list:  # 如果传入为id_list，则通过tokenizer转回文本
            while self.tokenizer.pad_token_id in label:  # 移除[PAD] token
                label.remove(self.tokenizer.pad_token_id)
            label = ''.join(self.tokenizer.convert_ids_to_tokens(label))  # 将id列表转换为文本
        if label not in self.label_dict:  # 如果标签不在字典中，抛出异常
            raise ValueError(f'Lable Error: "{label}" not in label_dict {list(self.label_dict)}.')
        sub_labels = self.label_dict[label]  # 获取对应的子标签列表
        ret = {'sub_labels': sub_labels}  # 构建返回结果
        token_ids = [_id[1:-1] for _id in self.tokenizer(sub_labels)['input_ids']]  # 转换子标签为token id
        for i in range(len(token_ids)):  # 对每个子标签的token id进行截断或补齐
            token_ids[i] = token_ids[i][:self.max_label_len]  # 截断超过最大长度的部分
            if len(token_ids[i]) < self.max_label_len:  # 补齐不足最大长度的部分
                token_ids[i] = token_ids[i] + [self.tokenizer.pad_token_id] * (self.max_label_len - len(token_ids[i]))
        ret['token_ids'] = token_ids  # 添加token ids到返回结果
        return ret  # 返回最终结果

    def batch_find_sub_labels(self, label: List[Union[list, str]]):
        """
        批量找到子标签。

        Args:
            label (List[Union[list, str]]): 标签列表, [[4510, 5554], [860, 5509]] or ['体育', '电脑']

        Returns:
            list -> [
                        {
                            'sub_labels': ['笔记本', '电脑'],
                            'token_ids': [[5011, 6381, 3315], [4510, 5554]]
                        },
                        ...
                    ]
        """
        return [self.find_sub_labels(l) for l in label]  # 对每个标签调用find_sub_labels方法

    def get_common_sub_str(self, str1: str, str2: str):
        """
        寻找最大公共子串。

        Args:
            str1 (str): 字符串1
            str2 (str): 字符串2

        Returns:
            tuple -> (最大公共子串, 最大公共子串长度)
        """
        lstr1, lstr2 = len(str1), len(str2)  # 获取两个字符串的长度
        record = [[0 for i in range(lstr2 + 1)] for j in range(lstr1 + 1)]  # 初始化动态规划矩阵
        p = 0  # 最长匹配对应在str1中的最后一位
        maxNum = 0  # 最长匹配长度
        for i in range(lstr1):  # 遍历str1的每个字符
            for j in range(lstr2):  # 遍历str2的每个字符
                if str1[i] == str2[j]:  # 如果字符相等，更新动态规划矩阵
                    record[i+1][j+1] = record[i][j] + 1
                    if record[i+1][j+1] > maxNum:  # 更新最长匹配长度和位置
                        maxNum = record[i+1][j+1]
                        p = i + 1
        return str1[p-maxNum:p], maxNum  # 返回最大公共子串及其长度

    def hard_mapping(self, sub_label: str):
        """
        强匹配函数，当模型生成的子label不存在时，通过最大公共子串找到重合度最高的主label。

        Args:
            sub_label (str): 子label。

        Returns:
            str: 主label。
        """
        label, max_overlap_str = '', 0  # 初始化主标签和最大重叠长度
        for main_label, sub_labels in self.label_dict.items():  # 遍历标签字典
            overlap_num = 0  # 初始化当前主标签的总重叠长度
            for s_label in sub_labels:  # 遍历子标签
                overlap_num += self.get_common_sub_str(sub_label, s_label)[1]  # 累加最大公共子串长度
            if overlap_num >= max_overlap_str:  # 更新最大重叠长度和主标签
                max_overlap_str = overlap_num
                label = main_label
        return label  # 返回最匹配的主标签

    def find_main_label(self, sub_label: List[Union[list, str]], hard_mapping=True):
        """
        通过子标签找到父标签。

        Args:
            sub_label (List[Union[list, str]]): 子标签, 文本型 或 id_list, e.g. -> '苹果' or [5741, 3362]
            hard_mapping (bool): 当生成的词语不存在时，是否一定要匹配到一个最相似的label。

        Returns:
            dict -> {
                'label': '水果',
                'token_ids': [3717, 3362]
            }
        """
        if type(sub_label) == list:  # 如果传入为id_list，则通过tokenizer转回文本
            pad_token_id = self.tokenizer.pad_token_id
            while pad_token_id in sub_label:  # 移除[PAD] token
                sub_label.remove(pad_token_id)
            sub_label = ''.join(self.tokenizer.convert_ids_to_tokens(sub_label))  # 将id列表转换为文本
        main_label = '无'  # 初始化主标签为“无”
        for label, s_labels in self.label_dict.items():  # 遍历标签字典
            if sub_label in s_labels:  # 如果子标签在子标签列表中，找到主标签
                main_label = label
                break
        if main_label == '无' and hard_mapping:  # 如果未找到主标签且开启强匹配
            main_label = self.hard_mapping(sub_label)  # 使用硬匹配找到最相似的主标签
        ret = {
            'label': main_label,
            'token_ids': self.tokenizer(main_label)['input_ids'][1:-1]  # 转换主标签为token id
        }
        return ret  # 返回最终结果

    def batch_find_main_label(self, sub_label: List[Union[list, str]], hard_mapping=True):
        """
        批量通过子标签找父标签。

        Args:
            sub_label (List[Union[list, str]]): 子标签列表, ['苹果', ...] or [[5741, 3362], ...]

        Returns:
            list: [
                    {
                    'label': '水果',
                    'token_ids': [3717, 3362]
                    },
                    ...
            ]
        """
        return [self.find_main_label(l, hard_mapping) for l in sub_label]  # 对每个子标签调用find_main_label方法


if __name__ == '__main__':
    from rich import print  # 导入 rich 库，用于美化打印输出
    from transformers import AutoTokenizer  # 导入 Hugging Face 的分词器

    tokenizer = AutoTokenizer.from_pretrained(pc.pre_model)  # 加载预训练分词器
    verbalizer = Verbalizer(
        verbalizer_file=pc.verbalizer,  # 加载verbalizer文件
        tokenizer=tokenizer,  # 使用分词器
        max_label_len=2  # 设置最大标签长度
    )
    print(verbalizer.label_dict)  # 打印加载的标签字典
    # 测试单个标签
    label = [4510, 5554]  # 输入的标签可以是id列表或文本
    ret = verbalizer.find_sub_labels(label)  # 查找子标签
    print(ret)  # 打印结果
#